import cx_Oracle
import pandas as pd
from typing import List, Dict
from vanna.remote import VannaDefault
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.milvus.milvus_vector import Milvus_VectorStore
from openai import OpenAI
import os
import numpy as np
import dashscope
from dashscope import TextEmbedding

# 1. 首先定义向量化类
class DashScopeEmbedding:
    def __init__(self):
        self.api_key = os.getenv("DASHSCOPE_API_KEY")
        dashscope.api_key = self.api_key
        self.dimension = 1536  # text-embedding-v2 的输出维度
        
    def encode_documents(self, texts):
        if isinstance(texts, str):
            texts = [texts]
            
        embeddings = []
        for text in texts:
            try:
                resp = TextEmbedding.call(
                    model='text-embedding-v2',
                    input=text,
                    api_key=self.api_key
                )
                if resp.status_code == 200:
                    # 处理字典格式的 embedding
                    embedding_data = resp.output['embeddings'][0]
                    if isinstance(embedding_data, dict):
                        # 如果是字典，尝试获取 embedding 值
                        if 'embedding' in embedding_data:
                            embedding = np.array(embedding_data['embedding'], dtype=np.float32)
                        else:
                            print(f"Embedding dict keys: {embedding_data.keys()}")
                            embedding = np.zeros(self.dimension, dtype=np.float32)
                    else:
                        embedding = np.array(embedding_data, dtype=np.float32)
                    
                    embeddings.append(embedding)
                else:
                    print(f"Error getting embedding: {resp.code} - {resp.message}")
                    embeddings.append(np.zeros(self.dimension, dtype=np.float32))
            except Exception as e:
                print(f"Exception during embedding: {str(e)}")
                print(f"Response content: {resp.output if 'resp' in locals() else 'No response'}")
                embeddings.append(np.zeros(self.dimension, dtype=np.float32))
        
        # 确保所有向量都是正确的形状
        embeddings = [e.reshape(self.dimension) if e.size == self.dimension else np.zeros(self.dimension, dtype=np.float32) for e in embeddings]
        return np.vstack(embeddings)
    
    def encode_queries(self, texts):
        return self.encode_documents(texts)
        
    def __call__(self, text):
        embeddings = self.encode_documents(text)
        return embeddings if len(embeddings.shape) > 1 else embeddings[0]

# 2. 初始化 OpenAI 客户端
client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    default_headers = {"x-foo": "true"}
)

# 3. 定义 MyVanna 类
class MyVanna(Milvus_VectorStore, OpenAI_Chat):
    def __init__(self, client=None, config=None):
        config = config or {}
        # 使用自定义的向量化函数
        config['embedding_function'] = DashScopeEmbedding()
        Milvus_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, client=client, config=config)

# 4. 定义初始化函数
def init_vanna(username: str, password: str, host: str, port: str, service_name: str):
    """
    初始化并配置 Vanna 实例，使用 Milvus 向量存储
    """
    vn = MyVanna(
        client=client, 
        config={
            "model": "qwen2.5-7b-instruct",
            "milvus_host": "192.168.2.110",  # Milvus 服务器地址
            "milvus_port": "19530",          # Milvus 服务端口
            "collection_name": "sql_vectors", # 集合名称
            "use_local_models": False         # 禁用本地模型下载
        }
    )
    vn.max_tokens = 800
    vn.temperature = 0.5
    dsn = cx_Oracle.makedsn(host, port, service_name=service_name)
    vn.connect_to_oracle(
        dsn=dsn, 
        user=username, 
        password=password
    )
    return vn

def train_vanna_with_ddl(vn, ddl_list: List[Dict]) -> List[str]:
    """
    使用 DDL 语句训练 Vanna
    
    Args:
        vn: Vanna 实例
        ddl_list: 包含 DDL 语句和描���的字典列表
    
    Returns:
        List[str]: 训练记录的 ID 列表
    """
    training_ids = []
    for ddl_info in ddl_list:
        try:
            # 使用 Vanna 的 train 方法进行训练
            training_id = vn.train(    # DDL 训练不需要 SQL
                ddl=ddl_info['ddl']
            )
            training_ids.append(training_id)
            print(f"成功训练 DDL，训练 ID: {training_id}")
            print(f"DDL: {ddl_info['ddl'][:100]}...")  # 只打印前100个字符
            print(f"描述: {ddl_info['description']}")
            print("-" * 80)
        except Exception as e:
            print(f"训练失败: {str(e)}")
            print(f"DDL: {ddl_info['ddl'][:100]}...")
            print("-" * 80)
    
    return training_ids

def connect_to_oracle(username: str, password: str, host: str, port: str, service_name: str) -> cx_Oracle.Connection:
    """
    连接到Oracle数据库
    """
    dsn = cx_Oracle.makedsn(host, port, service_name=service_name)
    connection = cx_Oracle.connect(username, password, dsn)
    return connection

def get_table_metadata(connection: cx_Oracle.Connection) -> pd.DataFrame:
    """
    获取所有表的元数据信息，包括表名和表注释，
    排除备份表（表名包含BAK或以数字结尾的表）
    """
    query = """
    SELECT 
        t.owner as OWNER,
        t.table_name as TABLE_NAME,
        t.comments as COMMENTS
    FROM all_tab_comments t
    WHERE t.owner = 'RESM_WHJT_PROD'
    AND t.table_type = 'TABLE'
    AND t.table_name NOT LIKE '%BAK%'  -- 排除包含BAK的表
    AND NOT REGEXP_LIKE(t.table_name, '[0-9]$')  -- 排除以数字结尾的表
    AND NOT REGEXP_LIKE(t.table_name, '[0-9]{8}')  -- 排除包含8位数字(日期)的表
    """
    
    return pd.read_sql(query, connection)

def get_column_metadata(connection: cx_Oracle.Connection) -> pd.DataFrame:
    """
    获取所有列的元数据信息，包括列名、数据类型和注释，
    排除备份表（表名包含BAK或以数字结尾的表）
    """
    query = """
    SELECT 
        c.owner as OWNER,
        c.table_name as TABLE_NAME,
        c.column_name as COLUMN_NAME,
        c.data_type as DATA_TYPE,
        c.data_length as DATA_LENGTH,
        c.nullable as NULLABLE,
        cc.comments as COMMENTS
    FROM all_tab_columns c
    LEFT JOIN all_col_comments cc 
        ON c.owner = cc.owner 
        AND c.table_name = cc.table_name 
        AND c.column_name = cc.column_name
    WHERE c.owner = 'RESM_WHJT_PROD'
    AND c.table_name NOT LIKE '%BAK%'  -- 排除包含BAK的表
    AND NOT REGEXP_LIKE(c.table_name, '[0-9]$')  -- 排除以数字结尾的表
    AND NOT REGEXP_LIKE(c.table_name, '[0-9]{8}')  -- 排除包含8位数字(日期)的表
    ORDER BY c.owner, c.table_name, c.column_id
    """
    
    return pd.read_sql(query, connection)

def generate_ddl_statements(tables_df: pd.DataFrame, columns_df: pd.DataFrame) -> List[Dict]:
    """
    生成用于Vanna训练的DDL语句和描述
    """
    ddl_list = []
    
    for _, table in tables_df.iterrows():
        table_columns = columns_df[
            (columns_df['OWNER'] == table['OWNER']) & 
            (columns_df['TABLE_NAME'] == table['TABLE_NAME'])
        ]
        
        # 构建CREATE TABLE语句
        ddl = f"CREATE TABLE {table['OWNER']}.{table['TABLE_NAME']} (\n"
        column_definitions = []
        
        for _, column in table_columns.iterrows():
            col_def = f"    {column['COLUMN_NAME']} {column['DATA_TYPE']}"
            if column['DATA_TYPE'] in ('VARCHAR2', 'CHAR'):
                col_def += f"({column['DATA_LENGTH']})"
            if column['NULLABLE'] == 'N':
                col_def += " NOT NULL"
            if column['COMMENTS']:
                col_def += f" -- {column['COMMENTS']}"
            column_definitions.append(col_def)
        
        ddl += ",\n".join(column_definitions)
        ddl += "\n)"
        
        # 如果有表注释，添加注释语句
        if table['COMMENTS']:
            ddl += f"\n\nCOMMENT ON TABLE {table['OWNER']}.{table['TABLE_NAME']} IS '{table['COMMENTS']}'"
        
        # 为Vanna准备训练数据
        ddl_info = {
            'ddl': ddl,
            'description': f"表 {table['OWNER']}.{table['TABLE_NAME']} 的结构定义。"
            f"该表{table['COMMENTS'] if table['COMMENTS'] else '未提供描述'}。"
        }
        ddl_list.append(ddl_info)
    
    return ddl_list

def generate_example_queries(tables_df: pd.DataFrame, columns_df: pd.DataFrame) -> List[Dict]:
    """
    为每个表生成示例查询语句和描述
    """
    query_list = []
    
    for _, table in tables_df.iterrows():
        table_columns = columns_df[
            (columns_df['OWNER'] == table['OWNER']) & 
            (columns_df['TABLE_NAME'] == table['TABLE_NAME'])
        ]
        
        # 构建注释说明
        table_description = f"查询{table['COMMENTS'] if table['COMMENTS'] else table['TABLE_NAME']}"
        
        # 构建SELECT语句
        query = f"-- {table_description}\nSELECT\n"
        column_definitions = []
        
        for _, column in table_columns.iterrows():
            # 构建列定义，包含注释
            col_def = column['COLUMN_NAME']
            if column['COMMENTS']:
                col_def += f", -- {column['COMMENTS']}"
            column_definitions.append(col_def)
        
        query += "\n".join(column_definitions)
        query += f"\nFROM {table['OWNER']}.{table['TABLE_NAME']}"
        
        # 准备训练数据
        query_info = {
            'question': table_description,
            'sql': query,
            'documentation': f"""
这个查询返回{table['COMMENTS'] if table['COMMENTS'] else table['TABLE_NAME']}的所有信息。
包含以下字段：
{', '.join([f"{col['COLUMN_NAME']}({col['COMMENTS'] if col['COMMENTS'] else '无说明'})" 
           for _, col in table_columns.iterrows()])}
"""
        }
        query_list.append(query_info)
    
    return query_list

def train_vanna_with_queries(vn, query_list: List[Dict]) -> List[str]:
    """
    使用生成的查询语句训练 Vanna
    """
    training_ids = []
    for query_info in query_list:
        try:
            training_id = vn.train(
                question=query_info['question'] + query_info['documentation'],
                sql=query_info['sql']
            )
            training_ids.append(training_id)
            print(f"成功训练查询，训练 ID: {training_id}")
            print(f"问题: {query_info['question']}{query_info['documentation']}")
            print(f"SQL: \n{query_info['sql']}")
            print("-" * 80)
        except Exception as e:
            print(f"训练失败: {str(e)}")
            print(f"SQL: {query_info['sql'][:100]}...")
            print("-" * 80)
    
    return training_ids

def main():
    # 数据库连接配置
    config = {
        'username': 'RESM_WHJT_PROD',
        'password': 'rZzj82q5Q2QBZ765zR6sU',
        'host': '192.168.2.105',
        'port': '1521',
        'service_name': 'ORCLPDB1'
    }
    
    try:
        # 连接数据库
        connection = connect_to_oracle(**config)
        
        # 获取元数据
        tables_df = get_table_metadata(connection)
        columns_df = get_column_metadata(connection)
        
        print(f"获取到 {len(tables_df)} 个表的元数据")
        
        # 初始化 Vanna
        vn = init_vanna(**config)
        
        # 1. 首先训练 DDL
        print("开始训练表结构(DDL)...")
        ddl_list = generate_ddl_statements(tables_df, columns_df)
        print(f"生成了 {len(ddl_list)} 个 DDL 语句")
        ddl_training_ids = train_vanna_with_ddl(vn, ddl_list)
        print(f"DDL训练完成，共训练 {len(ddl_training_ids)} 个表结构")
        
        # 2. 然后训练示例查询
        print("\n开始训练示例查询...")
        query_list = generate_example_queries(tables_df, columns_df)
        print(f"生成了 {len(query_list)} 个查询语句")
        query_training_ids = train_vanna_with_queries(vn, query_list)
        print(f"查询训练完成，共训练 {len(query_training_ids)} 个查询")
        
        # 获取训练数据总数
        training_data = vn.get_training_data()
        print("\n训练总结:")
        print(f"成功训练 {len(ddl_training_ids)} 个表结构")
        print(f"成功训练 {len(query_training_ids)} 个示例查询")
        print(f"Milvus 中总共有 {len(training_data)} 条训练数据")
        
    except cx_Oracle.Error as error:
        print(f"数据库错误: {error}")
    except Exception as e:
        print(f"其他错误: {str(e)}")
        import traceback
        traceback.print_exc()
    finally:
        if 'connection' in locals():
            connection.close()

if __name__ == "__main__":
    main()
